import cv2
import numpy as np
from matplotlib import pyplot as plt


def guided_filter(I, p, r, eps):
    """
    执行引导滤波。
    参数:
    - I: 引导图像（灰度图/单通道）
    - p: 输入滤波图像（灰度图/单通道）
    - r: 局部窗口半径
    - eps: 正则化参数
    """
    if len(p.shape) > 2:
        p = p[:, :, 0]  # 假设p错误地变成了三维，使用第一个通道

    p3 = np.repeat(p[:, :, np.newaxis], 3, axis=2)
    mean_I = cv2.boxFilter(I, cv2.CV_64F, (r, r))
    mean_p = cv2.boxFilter(p3, cv2.CV_64F, (r, r))
    corr_I = cv2.boxFilter(I * I, cv2.CV_64F, (r, r))
    corr_Ip = cv2.boxFilter(I * p3, cv2.CV_64F, (r, r))
    # 2
    var_I = corr_I - mean_I * mean_I
    cov_Ip = corr_Ip - mean_I * mean_p
    # 3
    a = cov_Ip / (var_I + eps)
    b = mean_p - a * mean_I
    # 4
    mean_a = cv2.boxFilter(a, cv2.CV_64F, (r, r))
    mean_b = cv2.boxFilter(b, cv2.CV_64F, (r, r))
    # 5
    q = mean_a * I + mean_b
    return q

def dark_channel(img, size=7):
    b,g,r = cv2.split(img)
    min_img = cv2.min(r, cv2.min(g, b))  # 取最暗通道
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (size, size))
    dc_img = cv2.erode(min_img, kernel)
    return dc_img

def get_atmo(img, top_percent=0.001):
    # 计算暗通道
    dark_img=dark_channel(img)
    flat_dark_img = dark_img.ravel()
    flat_img = img.reshape(-1, 3)

    # 选择暗通道图像中最亮的 top_percent 个像素
    num_pixels = int(top_percent * len(flat_dark_img))
    indices = np.argpartition(flat_dark_img, -num_pixels)[-num_pixels:]

    # 从这些像素中选取原始图像中最大RGB值的像素作为大气光A
    A = np.max(flat_img[indices], axis=0)

    return A

def get_trans(img, A, omega=0.95, t0=0.1):
    # 计算暗通道


    # 计算透射率，使用大气光的最小值来计算透射率
    transmission = 1 - omega * dark_channel(img / A.min())

    # 使用引导滤波器对透射率进行平滑
    t = guided_filter(img, transmission, r=23, eps=1e-3)

    # 设置透射率的最小值，避免过暗的结果
    t = np.maximum(t, t0)

    return t

def recover(A, img, t, t0):
    img = np.asarray(img)  # 确保 img 是数组
    if img.ndim != 3:
        raise ValueError("输入图像需要是三维的")

    t = np.clip(t, t0, 1)  # 限制透射率的最小值，防止图像过暗

    res = np.empty(img.shape, img.dtype)  # 初始化输出图像

    for index in range(3):  # 对于 RGB 的每个通道
        res[:, :, index] = (img[:, :, index] - A[index]) / t[:, :, index] + A[index]

    return res

def normalize_image(image):
    image_min = image.min()
    image_max = image.max()
    return (image - image_min) / (image_max - image_min)
def ret_dehaze(img):
    img = img/255
    im_dark = dark_channel(img)
    A = get_atmo(img)
    trans = get_trans(img, A)
    trans = np.clip(trans, 0.0, 1.0)
    result = recover(A, img, trans, 0.1)
    result = normalize_image(result)
    dehazed_img_8bit = (result * 255).astype(np.uint8)

    return dehazed_img_8bit


# if __name__ == '__main__':
#     img = cv2.imread('images_origin/1.png') / 255
#
#     plt.imshow(img[:, :, ::-1])
#     plt.show()
#
#     im_dark = dark_channel(img)
#     plt.imshow(im_dark, 'gray')
#     plt.show()
#
#     A = get_atmo(img)
#     trans = get_trans( img,A)
#     trans = np.clip(trans, 0.0, 1.0)
#     plt.imshow(trans, 'gray')
#     plt.show()
#
#     result = recover(A, img, trans,0.1)
#     result = normalize_image(result)
#     plt.imshow(result, 'gray')
#     plt.show()
